"""
Examples demonstrating the new tool system with Pydantic integration
"""
import os
from typing import Literal, List, Dict, Optional
from pydantic import BaseModel, Field
from sonetto.agents.user import UserProxy
from sonetto.agents.tool import ToolAgent
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset, field, validate
from sonetto.core.message import Message

def main():
    # 获取API配置
    api_key = os.getenv("API_KEY")
    if not api_key:
        raise ValueError("请设置API_KEY环境变量")
    
    base_url = os.getenv("API_BASE_URL")

    # 创建工具集
    toolset = McpToolset()

    # 方式1：使用field装饰器定义参数
    @toolset.register_tool(
        description="获取指定城市的天气信息"
    )
    def get_weather(
        city: str = field(
            description="城市名称",
            examples=["北京", "上海", "广州"]
        ),
        unit: Literal["celsius", "fahrenheit"] = field(
            default="celsius",
            description="温度单位",
            choices=["celsius", "fahrenheit"]
        )
    ) -> str:
        temperature = 25 if unit == "celsius" else 77
        return f"{city}的天气是晴朗，温度{temperature}度{unit}"

    # 方式2：使用validate装饰器添加约束
    @toolset.register_tool(
        description="计算数学表达式"
    )
    @validate("precision", min_value=0, max_value=10)
    def calculate(
        expression: str,
        precision: int = 2
    ) -> float:
        result = eval(expression)
        return round(float(result), precision)

    # 方式3：使用Pydantic模型定义参数
    class SearchParameters(BaseModel):
        query: str = Field(
            description="搜索关键词",
            min_length=1,
            max_length=100
        )
        limit: int = Field(
            default=10,
            description="返回结果数量",
            ge=1,
            le=100
        )
        category: Optional[str] = Field(
            default=None,
            description="搜索类别",
            examples=["news", "books", "articles"]
        )

    @toolset.register_tool(
        description="搜索数据库",
        model=SearchParameters
    )
    def search_database(params: SearchParameters) -> List[Dict]:
        # 模拟搜索结果
        results = [
            {"title": f"Result {i}", "score": 0.9 - i * 0.1}
            for i in range(params.limit)
        ]
        if params.category:
            results = [r for r in results if r["score"] > 0.5]
        return results

    # 创建调度器和代理
    scheduler = Scheduler()
    
    user = UserProxy()
    scheduler.register_agent(user)

    tool_agent = ToolAgent(
        agent_id="tool_assistant",
        api_key=api_key,
        base_url=base_url,
        toolset=toolset,
        temperature=0.7
    )
    scheduler.register_agent(tool_agent)

    # 设置消息流
    scheduler.set_agent_order(["user", "tool_assistant", "user"])

    print("=== Sonetto 工具系统示例 ===")
    print("可用工具：")
    for tool in toolset.list_tools():
        print(f"- {tool}")
    print("\n输入 'exit' 退出")
    print("===========================")

    while True:
        try:
            user_input = input("\n你: ").strip()
            if user_input.lower() == 'exit':
                break

            # 创建初始消息
            message = Message(
                content=user_input,
                sender=user.agent_id,
                receiver=tool_agent.agent_id
            )

            # 处理消息
            response = scheduler.step(message)
            while response:
                response = scheduler.step(response)

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main()
